#pragma once
#include <cassert>
#include <iostream>

namespace ns_RBTree
{
    enum COLOR
    {
        BLACK,
        RED
    };

    template <class T>
    struct RBTreeNode
    {
    public:
        T _key;
        COLOR _clr;
        RBTreeNode<T> *_left;
        RBTreeNode<T> *_right;
        RBTreeNode<T> *_parent;

        RBTreeNode(const T &key = T()) : _key(key), _clr(RED), _left(nullptr), _right(nullptr), _parent(nullptr)
        {
        }
    };

    template <class T>
    struct RBTree_Iterator
    {
        typedef RBTreeNode<T> Node;

        Node *_root;

        RBTree_Iterator(Node *root = nullptr) : _root(root)
        {
        }

        T &operator*()
        {
            return _root->_key;
        }

        T *operator->()
        {
            return &_root->_key;
        }

        RBTree_Iterator &operator++()
        {
            if (_root->_right)
            {
                Node *cur = _root->_right;
                while (cur->_left)
                {
                    cur = cur->_left;
                }
                _root = cur;
            }
            else
            {
                Node *parent = _root->_parent;
                Node *cur = _root;
                while (parent && parent->_right == cur)
                {
                    cur = cur->_parent;
                    parent = parent->_parent;
                }
                _root = parent;
            }

            return *this;
        }

        RBTree_Iterator &operator--()
        {
            if (_root->_left)
            {
                Node *cur = _root->_left;
                while (cur->_right)
                {
                    cur = cur->_right;
                }
                _root = cur;
            }
            else
            {
                Node *parent = _root->_parent;
                Node *cur = _root;
                while (parent && parent->_left == cur)
                {
                    cur = parent;
                    parent = parent->_parent;
                }
                _root = parent;
            }

            return *this;
        }

        bool operator==(const RBTree_Iterator &rb) const
        {
            return _root == rb._root;
        }

        bool operator!=(const RBTree_Iterator &rb) const
        {
            return _root != rb._root;
        }
    };

    template <class K, class T, class KeyOfT>
    class RBTree
    {
    public:
        typedef RBTreeNode<T> Node;
        typedef RBTree_Iterator<T> iterator;

        iterator begin()
        {
            Node *Left = _root;
            while (Left->_left)
            {
                Left = Left->_left;
            }

            return iterator(Left);
        }

        iterator end()
        {
            return iterator(nullptr);
        }

        RBTree() : _root(nullptr)
        {
        }

        RBTree(const RBTree<K, T, KeyOfT> &rb)
        {
            _root = copy(rb._root);
        }

        void swap(Node *&a, Node *&b)
        {
            Node *tmp = a;
            a = b;
            b = tmp;
        }

        RBTree<K, T, KeyOfT> &operator=(RBTree<K, T, KeyOfT> rb)
        {
            swap(_root, rb._root);

            return *this;
        }

        Node *copy(Node *rootB)
        {
            if (rootB == nullptr)
                return nullptr;
            Node *root = new Node(rootB->_key);
            root->_clr = rootB->_clr;

            root->_left = copy(rootB->_left);
            root->_right = copy(rootB->_right);

            if (root->_left)
                root->_left->_parent = root;
            if (root->_right)
                root->_right->_parent = root;

            return root;
        }

        std::pair<iterator, bool> insert(const T &key)
        {
            if (_root == nullptr)
            {
                _root = new Node(key);
                _root->_clr = BLACK;
                return std::make_pair(iterator(_root), true);
            }
            Node *cur = _root;
            Node *parent = nullptr;
            KeyOfT kot;
            while (cur)
            {
                if (kot(cur->_key) > kot(key))
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else if (kot(cur->_key) < kot(key))
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else
                    return std::make_pair(iterator(cur), false);
            }
            cur = new Node(key);
            Node *newnode = cur;
            if (kot(parent->_key) > kot(key))
                parent->_left = cur;
            else
                parent->_right = cur;
            cur->_parent = parent;

            // 调整平衡
            while (parent && parent->_clr == RED)
            {
                Node *grandParent = parent->_parent;
                assert(grandParent);
                if (grandParent->_left == parent)
                {
                    Node *uncle = grandParent->_right;
                    if (uncle && uncle->_clr == RED)
                    {
                        parent->_clr = BLACK;
                        uncle->_clr = BLACK;
                        grandParent->_clr = RED;

                        cur = grandParent;
                        parent = cur->_parent;
                    }
                    else
                    {
                        if (parent->_left == cur)
                        {
                            RotateR(grandParent);
                            grandParent->_clr = RED;
                            parent->_clr = BLACK;
                        }
                        else
                        {
                            RotateL(parent);
                            RotateR(grandParent);
                            grandParent->_clr = RED;
                            cur->_clr = BLACK;
                        }
                        break;
                    }
                }
                else
                {
                    Node *uncle = grandParent->_left;
                    if (uncle && uncle->_clr == RED)
                    {
                        parent->_clr = BLACK;
                        uncle->_clr = BLACK;
                        grandParent->_clr = RED;

                        cur = grandParent;
                        parent = cur->_parent;
                    }
                    else
                    {
                        if (parent->_right == cur)
                        {
                            RotateL(grandParent);
                            grandParent->_clr = RED;
                            parent->_clr = BLACK;
                        }
                        else
                        {
                            RotateR(parent);
                            RotateL(grandParent);
                            grandParent->_clr = RED;
                            cur->_clr = BLACK;
                        }
                        break;
                    }
                }
            }
            _root->_clr = BLACK;

            return std::make_pair(iterator(newnode), true);
        }

        Node *Find(const K &key)
        {
            Node *cur = _root;
            while (cur)
            {
                if (cur->_key.first > key)
                    cur = cur->_left;
                else if (cur->_key.first < key)
                    cur = cur->_right;
                else
                    return iterator(cur);
            }
            return iterator(nullptr);
        }

        ~RBTree()
        {
            Destroy(_root);
            _root = nullptr;
        }

        void Destroy(Node *&root)
        {
            if (nullptr == root)
                return;
            Destroy(root->_left);
            Destroy(root->_right);
            delete root;
        }

        void Inorder()
        {
            _Inorder(_root);
            std::cout << std::endl;
        }

        void _Inorder(Node *root)
        {
            if (root == nullptr)
                return;
            _Inorder(root->_left);
            std::cout << root->_key << " ";
            _Inorder(root->_right);
        }

    private:
        void RotateR(Node *&root)
        {
            Node *Left = root->_left;
            Node *LeftR = Left->_right;

            root->_left = LeftR;
            if (LeftR)
                LeftR->_parent = root;

            Left->_right = root;
            if (root->_parent == nullptr)
            {
                _root = Left;
                _root->_parent = nullptr;
            }
            else
            {
                if (root->_parent->_left == root)
                    root->_parent->_left = Left;
                else
                    root->_parent->_right = Left;
                Left->_parent = root->_parent;
            }
            root->_parent = Left;
        }

        void RotateL(Node *&root)
        {
            Node *Right = root->_right;
            Node *RightL = Right->_left;

            root->_right = RightL;
            if (RightL)
                RightL->_parent = root;

            Right->_left = root;
            if (root->_parent == nullptr)
            {
                _root = Right;
                _root->_parent = nullptr;
            }
            else
            {
                if (root->_parent->_left == root)
                    root->_parent->_left = Right;
                else
                    root->_parent->_right = Right;
                Right->_parent = root->_parent;
            }
            root->_parent = Right;
        }

    private:
        Node *_root;
    };
}